import os.path
from typing import Union, Callable, List, Dict

from scripts.Log import Log
from scripts.device.AdbShell import PmListUserInfo
from scripts.util.SettingsHelper import SettingsHelper
from scripts.actions.Action import Action
from scripts.actions.Context import Context


class BackupUserAction(Action):
    KEY = Action.register('user')
    DISPLAY_NAME = 'Created users'

    FILE_NAME = 'user.json'

    def __init__(self, context: Context):
        super().__init__(BackupUserAction.KEY, BackupUserAction.DISPLAY_NAME, context)

    def _on_backup(self, progress: Union[None, Callable[[int, int, str], None]]):
        device = self.wait_for_device()
        with device.shell() as shell:
            progress(0, 1, 'backup user info')
            settings = SettingsHelper(os.path.join(self.data_dir(), BackupUserAction.FILE_NAME))

            # backup user id, name, flags
            users = shell.pm_list_users()
            users_info = [{'id': user.user_id, 'name': user.user_name, 'flags': user.flags} for user in users]
            settings.put('users', users_info)
            settings.save()
            Log.d('backup user: %s' % users_info)

            progress(1, 1, 'Done [%s]' % self.duration_time_str())

            # user profiles will be backup by system data backup. data under:
            #    /data/system/users
            #    /data/system_*
            #    /data/misc/user

    def _on_restore(self, progress: Union[None, Callable[[int, int, str], None]]):
        device = self.wait_for_device()
        with device.shell() as shell:
            progress(0, 1, 'restore user info')
            settings = SettingsHelper(os.path.join(self.data_dir(), BackupUserAction.FILE_NAME))

            # load from json file
            users_info: List[Dict[str, Union[int, str]]] = settings.get('users')
            users = [PmListUserInfo(info['id'], info['name'], info['flags']) for info in users_info]
            Log.d('restore user: %s' % users_info)

            # get users to create & max user id
            current_user_list = {user.user_id for user in shell.pm_list_users()}
            restore_user_list = {user.user_id for user in users}
            Log.d('ignore create users: %s' % (current_user_list & restore_user_list))
            restore_user_list -= current_user_list
            users = list(filter(lambda user: user.user_id in restore_user_list, users))
            users = sorted(users, key=lambda user: user.user_id)

            # create user
            shell.wait_for_boot_complete()
            max_create_user = shell.pm_get_max_users() - 1
            created_user = set()
            to_create_user = {user.user_id for user in users}
            for user in users:
                while len(created_user) < max_create_user:
                    user_id = shell.pm_create_user(user.user_name)
                    if user_id is None:
                        Log.e('create user failed')
                        progress(1, 1, '!! create user failed [%s]' % self.duration_time_str())
                        return
                    created_user.add(user_id)
                    if user_id == user.user_id:
                        break

                # remove unnecessary users
                remove_user = created_user - to_create_user
                for user_id in remove_user:
                    shell.pm_remove_user(user_id)
                created_user -= remove_user

            progress(1, 1, 'Done [%s]' % self.duration_time_str())

            # user profiles will be restored by system data restore. data under:
            #    /data/system/users
            #    /data/system_*
            #    /data/misc/user
